
#include "ce_vulkan_object.h"
#include "ce_vulkan_manager.h"
#include "ce_vulkan_shader.h"
#include "ce_vulkan_context.h"
#include "ce_vulkan_pipeline.h"
#include "ce_vulkan_pass.h"

CE_VulkanObject::CE_VulkanObject(CE_VulkanContext *context, CE_VulkanObjectTypeE type)
{
    this->_context = context;
    this->_flag = 0;
    this->_color[0] = 0.0f;
    this->_color[1] = 0.0f;
    this->_color[2] = 0.0f;
    this->_color[3] = 1.0f;
    this->_type = type;
}

void CE_VulkanObject::setVertexs(std::vector<Vertex> vertices, std::vector<uint32_t> indices)
{
    // std::cout << "init vertex and index buffer start..." << std::endl;
    this->_cleanupBuffers(); // Clean up old buffers before creating new ones

    this->_vertexSize = vertices.size();
    this->_indexSize = indices.size();

    if (this->_vertexSize == 0 || this->_indexSize == 0)
        return; // Nothing to create
    // Vertices
    CE_VulkanManager::getInstance()->createBuffer(
        VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
        &this->_vertexBuffer,
        &this->_vertexMemory,
        this->_vertexSize * sizeof(Vertex),
        vertices.data());

    // Indices
    CE_VulkanManager::getInstance()->createBuffer(
        VK_BUFFER_USAGE_INDEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
        &this->_indexBuffer,
        &this->_indexMemory,
        this->_indexSize * sizeof(uint32_t),
        indices.data());

    this->_flag = 1;
}
void CE_VulkanObject::_cleanupBuffers()
{
    VkDevice device = _context->getVKDevice();
    vkDeviceWaitIdle(device);
    
    if (_vertexBuffer != VK_NULL_HANDLE)
    {
        vkDestroyBuffer(device, _vertexBuffer, nullptr);
        _vertexBuffer = VK_NULL_HANDLE;
    }
    if (_vertexMemory != VK_NULL_HANDLE)
    {
        vkFreeMemory(device, _vertexMemory, nullptr);
        _vertexMemory = VK_NULL_HANDLE;
    }
    if (_indexBuffer != VK_NULL_HANDLE)
    {
        vkDestroyBuffer(device, _indexBuffer, nullptr);
        _indexBuffer = VK_NULL_HANDLE;
    }
    if (_indexMemory != VK_NULL_HANDLE)
    {
        vkFreeMemory(device, _indexMemory, nullptr);
        _indexMemory = VK_NULL_HANDLE;
    }
}

void CE_VulkanObject::setRendererPipeline(CE_VulkanPipeline *pipeline)
{
    this->_pipeline = pipeline;
    this->_createFramebuffers();
    this->_flag = 1;
}
void CE_VulkanObject::_createFramebuffers()
{
    this->_cleanupFramebuffers();

    std::cout << "create framebuffers start..." << std::endl;
    this->_swapChainFramebuffers.resize(this->_context->getSwapChainImageViews().size());
    for (size_t i = 0; i < this->_context->getSwapChainImageViews().size(); i++)
    {
        VkImageView attachments[] = {this->_context->getSwapChainImageViews()[i]};
        VkFramebufferCreateInfo framebufferInfo = {};
        framebufferInfo.sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO;
        framebufferInfo.renderPass = this->_pipeline->getPass()->getRenderPass();
        framebufferInfo.attachmentCount = 1;                                  // 指定附着的个数
        framebufferInfo.pAttachments = attachments;                           // 渲染流程对象用于描述附着信息的pAttachment数组
        framebufferInfo.width = this->_context->getSwapChainExtent().width;   // width和height用于指定帧缓冲的大小
        framebufferInfo.height = this->_context->getSwapChainExtent().height; // 交换链图像都是单层，layers设置为1
        framebufferInfo.layers = 1;

        if (vkCreateFramebuffer(this->_context->getVKDevice(), &framebufferInfo, nullptr, &this->_swapChainFramebuffers[i]) != VK_SUCCESS)
        {
            throw std::runtime_error("Failed to create framebuffer!");
        }
    }
    std::cout << "create framebuffers success..." << std::endl;

    this->_commandBuffers.resize(this->_swapChainFramebuffers.size());

    VkCommandBufferAllocateInfo allocInfo{};
    allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    allocInfo.commandPool = this->_context->getCommandPool();
    allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    allocInfo.commandBufferCount = (uint32_t)this->_commandBuffers.size();

    if (vkAllocateCommandBuffers(this->_context->getVKDevice(), &allocInfo, this->_commandBuffers.data()) != VK_SUCCESS)
    {
        throw std::runtime_error("Failed to allocate command buffers!");
    }
}
void CE_VulkanObject::_cleanupFramebuffers()
{
    if (!this->_commandBuffers.empty())
    {
        vkFreeCommandBuffers(this->_context->getVKDevice(), this->_context->getCommandPool(), static_cast<uint32_t>(this->_commandBuffers.size()), this->_commandBuffers.data());
        this->_commandBuffers.clear();
    }
    for (auto framebuffer : this->_swapChainFramebuffers)
    {
        vkDestroyFramebuffer(this->_context->getVKDevice(), framebuffer, nullptr);
    }
    this->_swapChainFramebuffers.clear();
}

void CE_VulkanObject::setClear(bool clear)
{
    this->_isClear = clear;
    this->_flag = 1;
}
void CE_VulkanObject::setColor(float r, float g, float b, float a)
{
    this->_color[0] = r;
    this->_color[1] = g;
    this->_color[2] = b;
    this->_color[3] = a;
    this->_flag = 1;
}

void CE_VulkanObject::updateRender()
{
    if (this->_flag == 0)
        return;
    if (!this->_pipeline || this->_commandBuffers.empty() || this->_vertexBuffer == VK_NULL_HANDLE || this->_indexBuffer == VK_NULL_HANDLE)
        return;

    this->_flag = 0;

    for (size_t i = 0; i < this->_commandBuffers.size(); i++)
    {
        vkResetCommandBuffer(this->_commandBuffers[i], 0);

        VkCommandBufferBeginInfo beginInfo{};
        beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
        // beginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;

        if (vkBeginCommandBuffer(this->_commandBuffers[i], &beginInfo) != VK_SUCCESS)
        {
            throw std::runtime_error("Failed to begin recording command buffer!");
        }

        // 添加图像布局转换屏障--待确定有用没
        VkImageMemoryBarrier barrier{};
        barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
        barrier.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED;
        barrier.newLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
        barrier.image = this->_context->getSwapChainImages()[i]; // 使用当前交换链图像
        barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        barrier.subresourceRange.baseMipLevel = 0;
        barrier.subresourceRange.levelCount = 1;
        barrier.subresourceRange.baseArrayLayer = 0;
        barrier.subresourceRange.layerCount = 1;
        barrier.srcAccessMask = 0;
        barrier.dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;

        vkCmdPipelineBarrier(
            this->_commandBuffers[i],
            VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
            VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
            0, 0, nullptr, 0, nullptr, 1, &barrier);

        // 开始渲染pass
        VkRenderPassBeginInfo renderPassInfo{};
        renderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
        renderPassInfo.renderPass = this->_pipeline->getPass()->getRenderPass();
        renderPassInfo.framebuffer = this->_swapChainFramebuffers[i];
        renderPassInfo.renderArea.offset = {0, 0};
        renderPassInfo.renderArea.extent = this->_context->getSwapChainExtent();

        VkClearValue clearColor{};
        renderPassInfo.clearValueCount = 1;
        clearColor.color = {{0.0f, 0.0f, 0.0f, 0.0f}};
        renderPassInfo.pClearValues = &clearColor;

        vkCmdBeginRenderPass(this->_commandBuffers[i], &renderPassInfo, VK_SUBPASS_CONTENTS_INLINE);

        if (this->_isClear)
        {
            VkClearAttachment clearAtt{};
            clearAtt.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
            clearAtt.clearValue.color = {{0.0f, 0.0f, 0.0f, 0.0f}}; // 红色，Alpha=0（透明）
            clearAtt.colorAttachment = 0;                           // 第0个颜色附件

            VkClearRect clearRect{};
            clearRect.rect = renderPassInfo.renderArea; // 与 RenderPass 范围一致
            clearRect.baseArrayLayer = 0;
            clearRect.layerCount = 1;

            vkCmdClearAttachments(this->_commandBuffers[i], 1, &clearAtt, 1, &clearRect);
        }

        vkCmdBindPipeline(this->_commandBuffers[i], VK_PIPELINE_BIND_POINT_GRAPHICS, this->_pipeline->getPipeline());

        // 该函数用于提交绘制操作到指令缓冲
        // 参数1：需要执行的指令缓冲对象
        // 参数2：顶点缓冲数量
        // 参数3：用于实例渲染，为1时表示不进行实例渲染
        // 参数4：用于定义着色器变量gl_VertexIndex的值
        // 参数5：用于定义着色器变量gl_InstanceIndex的值

        // Render scene
        VkDeviceSize offsets[1] = {0};
        vkCmdBindVertexBuffers(this->_commandBuffers[i], 0, 1, &this->_vertexBuffer, offsets);
        vkCmdBindIndexBuffer(this->_commandBuffers[i], this->_indexBuffer, 0, VK_INDEX_TYPE_UINT32);

        // 推送常量
        PushConstants pushConstants{};
        pushConstants.defaultColor[0] = this->_color[0];
        pushConstants.defaultColor[1] = this->_color[1];
        pushConstants.defaultColor[2] = this->_color[2];
        pushConstants.defaultColor[3] = this->_color[3];
        vkCmdPushConstants(this->_commandBuffers[i], this->_pipeline->getPipelineLayout(), VK_SHADER_STAGE_FRAGMENT_BIT, 0, sizeof(PushConstants), &pushConstants);
        vkCmdDrawIndexed(
            this->_commandBuffers[i],
            this->_indexSize, // 只绘制3个索引（第一个三角形）
            1,                // 实例数 （2的话代表绘制2个实例，也就是绘制两次）
            0,                // 第一个索引的偏移量
            0,                // 顶点偏移
            0                 // 实例偏移
        );

        vkCmdEndRenderPass(this->_commandBuffers[i]);

        if (vkEndCommandBuffer(this->_commandBuffers[i]) != VK_SUCCESS)
        {
            throw std::runtime_error("Failed to record command buffer!");
        }
        // std::cout << "create command buffer success" << std::endl;
    }
}
VkCommandBuffer CE_VulkanObject::getCommandBuffer(uint32_t index)
{
    if (index >= this->_commandBuffers.size())
    {
        return VK_NULL_HANDLE;
    }
    return this->_commandBuffers[index];
}

CE_VulkanObject::~CE_VulkanObject()
{
    this->_cleanupBuffers();
    this->_cleanupFramebuffers();
}